#include <linux/bpf.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/ipv6.h>
#include <linux/skbuff.h>
#include <bpf/bpf_tracing.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_endian.h>

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __type(key, __u32);         // IPv4地址
    __type(value, __u8);        // 标记(1=监控)
    __uint(max_entries, 1024);  // 足够容纳500个IPv4
} ipv4_monitor_list SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __type(key, struct in6_addr);  // IPv6地址
    __type(value, __u8);           // 标记(1=监控)
    __uint(max_entries, 1024);     // 足够容纳500个IPv6
} ipv6_monitor_list SEC(".maps");

struct stat_value {
    struct bpf_spin_lock lock;
    __u64 value;
};

struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH); // 使用LRU统计map提高性能
    __type(key, __u32);                  // IPv4地址
    __type(value, struct stat_value);    // 字节计数
    __uint(max_entries, 1024);
} ipv4_stats SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH); // 使用LRU统计map提高性能
    __type(key, struct in6_addr);        // IPv6地址
    __type(value, struct stat_value);    // 字节计数
    __uint(max_entries, 1024);
} ipv6_stats SEC(".maps");

SEC("kprobe/__netif_receive_skb")
int trace_receive(struct pt_regs *ctx) {
    struct sk_buff *skb = (struct sk_buff *)PT_REGS_PARM1(ctx);
    void *data = (void *)(long)skb->head + skb->network_header;
    void *data_end = data + skb->len;
    struct ethhdr *eth = data;
    if (data + sizeof(*eth) > data_end)
        return 0;
    __u16 proto = eth->h_proto;
    // 处理IPv4
    if (proto == bpf_htons(ETH_P_IP)) {
        struct iphdr *ip = data + sizeof(*eth);
        if (data + sizeof(*eth) + sizeof(*ip) > data_end)
            return 0;
        // 快速查找监控列表
        __u8 *monitored = bpf_map_lookup_elem(&ipv4_monitor_list, &ip->daddr);
        if (!monitored || *monitored != 1)
            return 0;
        // 更新统计
        struct stat_value *val_ptr = bpf_map_lookup_elem(&ipv4_stats, &ip->daddr);
        if (val_ptr) {
            bpf_spin_lock(&val_ptr->lock);
            val_ptr->value += skb->len;
            bpf_spin_unlock(&val_ptr->lock);
        } else {
            struct stat_value init_val = { .value = skb->len, .lock = {} };
            bpf_map_update_elem(&ipv4_stats, &ip->daddr, &init_val, BPF_NOEXIST);
        }
    }
    // 处理IPv6
    else if (proto == bpf_htons(ETH_P_IPV6)) {
        struct ipv6hdr *ip6 = data + sizeof(*eth);
        if (data + sizeof(*eth) + sizeof(*ip6) > data_end)
            return 0;
        // 快速查找监控列表
        __u8 *monitored = bpf_map_lookup_elem(&ipv6_monitor_list, &ip6->daddr);
        if (!monitored || *monitored != 1)
            return 0;
        // 更新统计
        struct stat_value *val_ptr = bpf_map_lookup_elem(&ipv6_stats, &ip6->daddr);
        if (val_ptr) {
            bpf_spin_lock(&val_ptr->lock);
            val_ptr->value += skb->len;
            bpf_spin_unlock(&val_ptr->lock);
        } else {
            struct stat_value init_val = { .value = skb->len, .lock = {} };
            bpf_map_update_elem(&ipv6_stats, &ip6->daddr, &init_val, BPF_NOEXIST);
        }
    }
    return 0;
}

char _license[] SEC("license") = "GPL";
